#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <signal.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <pwd.h>
#include <readline/readline.h>
#include <readline/history.h>
#include "job.h"

#define MAXARGS 128
#define MAXLINE 8192 
#define PIDLEN  11

extern void yylex();

extern job_info *jobs;          /* header for job link */
volatile sig_atomic_t child;    /* lock used to let shell wait */
sigset_t oldmask;               /* saved mask, not mask SIGCHLD */ 
char *argvbuf[MAXARGS];         /* saved all the argv */
int argcnt;                     /* count for argv */
char lastdir[MAXLINE];          /* saved last dir */
int ttyfd;                      /* fd to /dev/tty */
char buf[MAXLINE];              /* used to lex */
char *myptr;                    /* used to lex */
char *mylim;                    /* used to lex */

typedef struct _pid_pgid {
    pid_t pid;
    pid_t pgid;
} pid_pgid;

pid_pgid pid_pgids[MAXARGS];    /* save all pid pgid group */
int pindex;                     /* count for pid pgid group */

pid_t get_pgid(pid_t pid)
{
    int i;

    for (i = 0; i < pindex; i++)
        if (pid_pgids[i].pid == pid)
            return pid_pgids[i].pgid;

    return 0;
}

void del_pgid(pid_t pid)
{
    int i, j;

    for (i = 0; i < pindex; i++)
        if (pid_pgids[i].pid == pid) 
            break;

    pid_pgids[i].pid = 0;
    pid_pgids[i].pgid = 0;
    while (i+1 < pindex) {
        pid_pgids[i] = pid_pgids[i+1];
        i ++;
    }
    pindex --;
}

void reset_argvbuf()
{
    int i;

    for (i = 0; i < argcnt; i++) {
        free(argvbuf[i]);
        argvbuf[i] = NULL;
    }
    argcnt = 0;
}

void parent_handler(int sig)
{
    int status;
    pid_t pid, pgid, tcpgid;
    job_info *job;

    if (sig == SIGTTOU)                     /* handler for SIGTTOU */
        tcsetpgrp(ttyfd, getpid()); 
    else if (sig == SIGCHLD) {
        while ((pid = waitpid(-1, &status, 
                    WNOHANG|WUNTRACED|WCONTINUED)) > 0) { 
            pgid = get_pgid(pid);
            tcpgid = tcgetpgrp(ttyfd);
            if (WIFEXITED(status)) {        /* handler for child exit */
                job = find_job(0, pgid);
                if (job != NULL) {
                    job->run --;
                    job->alive --;
                    if (job->alive == 0) {
                        del_job(pgid);
                        job = NULL;
                    }
                    else if (job->run == 0)
                        set_job(pgid, STOP, SET_NO);

                    del_pgid(pid);          /* del pgid in the pid-pgids */

                    if (pgid = tcpgid && (job == NULL || job->run == 0)) {
                        tcsetpgrp(ttyfd, getpid()); 
                        child = 1;          /* unlock shell */
                    }
                }
            }
            else if (WIFSIGNALED(status)) { /* handler for child term */
                job = find_job(0, pgid);
                if (job != NULL) {
                    job->run --;
                    job->alive --;
                    if (job->alive == 0) {
                        printf("\nJob [%d] %s terminate\n", job->jid, 
                            job->cmdline);
                        del_job(pgid);
                        job = NULL;
                    }
                    else if (job->run == 0)
                        set_job(pgid, STOP, SET_NO);
                    
                    del_pgid(pid);          /* del pgid in the pid-pgids */

                    if (pgid == tcpgid && (job == NULL || job->run == 0)) {
                        tcsetpgrp(ttyfd, getpid()); 
                        child = 1;          /* unlock shell */
                    }
                }
            }
            else if (WIFSTOPPED(status)) {  /* handler for child stop */
                job = find_job(0, pgid);
                if (job != NULL) {
                    job->run --;
                    if (job->run == 0) {
                        set_job(pgid, STOP, SET_NO);
                        printf("\nJob [%d] %s Stopped\n", job->jid, 
                            job->cmdline);
                        if (pgid == tcpgid) {
                            tcsetpgrp(ttyfd, getpid()); 
                            child = 1;          /* unlock shell */
                        }
                    }
                }
            }
            else if (WIFCONTINUED(status)) {/* handler for child cont */
                job = find_job(0, pgid);
                if (job != NULL) {
                    job->run ++;
                    if (job->run == 1) {        /* stop -> continue */
                        if (pgid == tcpgid)
                            child = 1;          /* unlock shell */
                        else
                            set_job(pgid, RUN, SET_BG); 
                    }
                }
            }
        }
    }
}

/* if first arg is a builtin command, run it and return true */
int builtin_command(int argc, char *argv[])
{   
    job_info *job;
    pid_t pgid;
    int jid, i;
    HIST_ENTRY **h;
    char *dir;
    char cwd[MAXLINE];

    if (!strcmp(argv[0], "history")) {
        h = history_list();
        if (h != NULL) {
            i = 0;
            while (h[i] != NULL) {
                printf("%d: %s\n", i, h[i]->line);
                i++;
            }
        }
        return 1;
    }

    /* quit command */
    if (!strcmp(argv[0], "quit") || !strcmp(argv[0], "exit")) {
        reset_argvbuf();
        free_jobs();
        if (argc > 1)
            i = atoi(argv[1]);
        else
            i = 0;
        exit(i);
    }

    if (!strcmp(argv[0], "cd")) {   /* cd command */
        if (argc == 1) {
            if ((dir = getenv("HOME")) == NULL) {
                perror("cd error");
                return 1;
            }
        }
        else if (argc == 2) {
            if (!strcmp(argv[1], "-")) {
                dir = lastdir;
            }
            else if (!strcmp(argv[1], "~")) {
                if ((dir = getenv("HOME")) == NULL) {
                    perror("cd error");
                    return 1;
                }
            }
            else
                dir = argv[1];
        }
        else {
            printf("usage: cd [dir]\n");
            return 1;
        }

        getcwd(cwd, MAXLINE);
        if (chdir(dir) < 0) {
            perror("cd error");
            return 1;
        }

        strncpy(lastdir, cwd, strlen(cwd)+1); 
        return 1;
    }

    if (!strcmp(argv[0], "jobs")) { /* jobs command, show all bg jobs */
        print_jobs();
        return 1;
    }

    if (!strcmp(argv[0], "bg")) {   /* bg command, continue in the bg */
        if (argc < 2)
            pgid = 0;
        else if (argv[1][0] == '%')
            jid = atoi(&argv[1][1]);
        else
            pgid = atoi(argv[1]);

        job = find_job(jid, pgid); 
        if (job == NULL)
            printf("No such process!\n");
        else if (strcmp(job->status, RUN)) {
            printf("Job [%d] %s %s\n", job->jid, job->cmdline, RUN);
            kill(-job->pgid, SIGCONT);
        }
        else
            printf("job [%d] %s has been %s in bg\n", job->jid, 
                job->cmdline, RUN);
        return 1;
    }

    if (!strcmp(argv[0], "fg")) {   /* fg command, continue in the fg */
        if (argc < 2)
            pgid = 0;
        else if (argv[1][0] == '%')
            jid = atoi(&argv[1][1]);
        else
            pgid = atoi(argv[1]);

        job = find_job(jid, pgid); 
        if (job == NULL)
            printf("No such process!\n");
        else {
            job = set_job(job->pgid, RUN, SET_FG);
            printf("Job [%d] %s %s\n", job->jid, job->cmdline, job->status);
            tcsetpgrp(ttyfd, job->pgid);/* set tcpgrp */
            kill(-job->pgid, SIGCONT);
            while (child == 0)          /* wait child */
                sigsuspend(&oldmask);
            child = 0;
        }
        return 1;
    }

    return 0;
}

void eval(char *cmdline)
{
    int i, j, len;
    char *cmd, *p;

    if (argcnt == 0)
        return;

    for (i=0, j=0; argvbuf[i] != NULL; i++) {
        if (!strcmp(argvbuf[i], ";")) {
            argvbuf[i] = NULL;
            /* malloc cmdline for pipe cmd */
            p = strchr(cmdline, ';'); 
            len = p - cmdline;
            cmd = malloc(len+1);
            strncpy(cmd, cmdline, len+1);
            cmd[len] = '\0';
            if (do_pipe_cmd(j, i, argvbuf+j, cmd)) {
                j = i+1;
                cmdline = p+1;
            }
            else
                return ;
        }
    }

    if (i == j)
        return ;

    len = strlen(cmdline);
    cmd = malloc(len+1);
    strncpy(cmd, cmdline, len+1);
    do_pipe_cmd(j, i, argvbuf+j, cmd);
}

/* return 1: more cmd need run, 0: no more cmd need run */
int do_pipe_cmd(int start, int end, char *argv[], char *cmdline)
{
    int i, j, prepipe, bg=0;
    int prefd[2], postfd[2];
    int len = end - start;

    /* check if '|' is at the last */
    if (!strcmp(argv[len-1], "|") || (len > 1 && !strcmp(argv[len-2], "|")
        && !strcmp(argv[len-1], "&"))) {
        printf("syntax error at |\n");
        return 0;
    }

    /* check if '&' is at the last */
    if (!strcmp(argv[len-1], "&") && end != argcnt) {
        printf("syntax error at &\n");
        return 0;
    }
    else if (!strcmp(argv[len-1], "&") && end == argcnt) {
        bg = 1;
        argv[len-1] = NULL;
        len --;
        if (len == 0) {
            printf("syntax error at &\n");
            return 0;
        }
    }

    prepipe = 0;
    for (i = 0, j = 0; i < len; i++) {
        if (!strcmp(argv[i], "|")) {
            argv[i] = NULL;
            pipe(postfd);
            if (prepipe) {
                if (!do_simple_cmd(i-j, argv+j, prepipe, postfd, 
                    bg, cmdline))
                    return 1;
            }
            else { 
                if (!do_simple_cmd(i-j, argv+j, 0, postfd, bg, cmdline))
                    return 1;
            }
            prefd[0] = postfd[0];
            prefd[1] = postfd[1];
            prepipe = 1;
            j = i+1;
        }
    }

    if (i == j)
        return 1;

    if (prepipe)
        do_simple_cmd(i-j, argv+j, prefd, 0, bg, cmdline);
    else
        do_simple_cmd(i-j, argv+j, 0, 0, bg, cmdline);

    return 1;
}

/* return -1 if error */
int predo_for_redirect(int argc, char *argv[], int *re)
{
    int i, redirect, read_fd, write_fd;
    char *in, *out;

    for (i = 0, redirect = 0; i < argc; i++) {
        if (!strcmp(argv[i], "<")) {
            redirect += 1;
            argv[i] = NULL;
            in = argv[i+1];
        }
        else if (!(strcmp(argv[i], ">"))) {
            redirect += 2;
            argv[i] = NULL;
            out = argv[i+1];
        }
    }

    /* need redirect */
    switch (redirect) {
        case 1:
            if ((read_fd = open(in, O_RDONLY)) < 0) {
                fprintf(stderr, "open %s error\n", in);
                return -1;
            }
            dup2(read_fd, STDIN_FILENO);
            break;
        case 2:
            if ((write_fd = open(out, O_WRONLY|O_CREAT|O_TRUNC,
                        S_IRUSR|S_IWUSR)) < 0) {
                fprintf(stderr, "open %s error\n", out);
                return -1;
            }
            dup2(write_fd, STDOUT_FILENO);
            break;
        case 3:
            if ((read_fd = open(in, O_RDONLY)) < 0) {
                fprintf(stderr, "open %s error\n", in);
                return -1;
            }
            dup2(read_fd, STDIN_FILENO);

            if ((write_fd = open(out, O_WRONLY|O_CREAT|O_TRUNC,
                        S_IRUSR|S_IWUSR)) < 0) {
                fprintf(stderr, "open %s error\n", out);
                return -1;
            }
            dup2(write_fd, STDOUT_FILENO);
            break;
        default:
            break;
    }

    if (re != NULL)
        *re = redirect;

    return 0;
}   

/* return 1: more cmd need run, 0: no more cmd need run */
int do_simple_cmd(int argc, char *argv[], int prefd[], int postfd[], 
    int bg, char *cmdline)
{
    int ret = 1;
    sigset_t mask;
    static int count = 0;
    static pid_t pgid = 0, tmppid = 0;
    pid_t pid;
    int redirect;
    job_info *job;
    pid_pgid ppg;
    int fd1[2], fd2[2];
    char pgrp[PIDLEN];

    /* for more than one process in the pipe chain */
    if (prefd == 0 && postfd != 0) {    /* the first process */
        pipe(fd1);
    }
    else if (prefd != 0) {              /* the left processes */
        pipe(fd1);
        pipe(fd2);
    }

    if (prefd == 0) {
        sigemptyset(&mask);
        sigaddset(&mask, SIGCHLD);
        sigprocmask(SIG_BLOCK, &mask, &oldmask);
    }

    if (!builtin_command(argc, argv)) { /* not builtin-command */
        if ((pid = fork()) == 0) {      /* child process job */
            
            /* only one process in pipe chain */
            if (prefd == 0 && postfd == 0)
                setpgid(0, 0);
            else if (prefd == 0) {      /* the first pipe process */  
                close(fd1[1]);
                setpgid(0, 0);          /* set the pid to gid */
                read(fd1[0], pgrp, PIDLEN);
            }
            else {                      /* set other process pgid */
                close(fd1[1]);
                close(fd2[0]);
                read(fd1[0], pgrp, PIDLEN);
                setpgid(0, atoi(pgrp));
                write(fd2[1], "c", 1);
            }

            signal(SIGINT, SIG_DFL);
            signal(SIGTSTP, SIG_DFL);
            signal(SIGQUIT, SIG_DFL);
            signal(SIGCHLD, SIG_DFL);
            signal(SIGTTOU, SIG_DFL);
            sigprocmask(SIG_SETMASK, &oldmask, NULL);

            if (predo_for_redirect(argc, argv, &redirect) < 0)
                exit(1);

            if (prefd != 0 && (redirect != 1 || redirect != 3)) {
                close(prefd[1]);
                if (prefd[0] != STDIN_FILENO) {
                    dup2(prefd[0], STDIN_FILENO);
                    close(prefd[0]);
                }
            }

            if (postfd != 0 && (redirect != 2 || redirect != 3)) {
                close(postfd[0]);
                if (postfd[1] != STDOUT_FILENO) {
                    dup2(postfd[1], STDOUT_FILENO);
                    close(postfd[1]);
                }
            }

            if (execvp(argv[0], argv) < 0) {
                printf("%s Command not found.\n", argv[0]);
                exit(0);
            }
        }

        /* add pipe chain process count */
        count ++;
    
        /* only one process in pipe chain */
        if (prefd == 0 && postfd == 0) {
            pgid = pid;
        }
        else if (prefd == 0) {
            close(fd1[0]);
            pgid = pid;

            /* let temp child add the pipe chain process group */
            if ((tmppid = fork()) == 0) {
                close(fd1[0]);
                close(fd1[1]);
                setpgid(0, pgid);
                if (postfd != 0)
                    close(postfd[1]);
                pause();
                exit(0);
            }    
            close(fd1[1]);               /* close the write fd */
        }
        else {
            close(fd1[0]);
            close(fd2[1]);
            sprintf(pgrp, "%d", pgid);
            write(fd1[0], pgrp, strlen(pgrp)+1);
            read(fd2[1], pgrp, PIDLEN);
            kill(tmppid, SIGKILL);      /* kill the temp child */
            tmppid = 0;                 /* clear tmppid for next pipe chain */
            close(fd1[1]);
            close(fd2[0]);
        }

        /* add pid-pgid to the global pid_pgids */
        ppg.pid = pid;
        ppg.pgid = pgid;
        pid_pgids[pindex++] = ppg;

        if (postfd == 0) {
            add_job(pgid, count, count, RUN, cmdline);
            count = 0;                  /* clear count for next pipe chain */

            if (!bg) {                  /* parent wait fg child */
                tcsetpgrp(ttyfd, pgid); /* set tcpgrp to fg */
                pgid = 0;               /* clear pgid for next pipe chain */
                while (child == 0)      /* wait child */
                    sigsuspend(&oldmask);
                child = 0;
            }
            else {
                job = find_job(0, pgid);
                printf("Job [%d] %d\n", job->jid, job->pgid);
                pgid = 0;               /* clear pgid for next pipe chain */
            }

            ret = 0;
        }
        else
            close(postfd[1]);           /* close write for the next read */
    }
    else
        ret = 0;

    if (ret == 0)
        sigprocmask(SIG_UNBLOCK, &mask, NULL);

    return ret;
}

void set_prompt(char *prompt)
{
    char host[MAXLINE];
    char cwd[MAXLINE];
    struct passwd *pwp;

    if (gethostname(host, MAXLINE) < 0)
        strncpy(host, "unknown", sizeof("unknown"));

    if (getcwd(cwd, MAXLINE) == 0) {
        strncpy(cwd, "unknown", sizeof("unknown"));
    }

    pwp = getpwuid(getuid());
    sprintf(prompt, "%s@%s:%s > ", pwp?pwp->pw_name:"unknown", host, cwd);
}

int main(int argc, char *argv[], char *envp[])
{
    char *cmdline;
    char prompt[MAXLINE];

    signal(SIGINT, SIG_IGN);
    signal(SIGTSTP, SIG_IGN);
    signal(SIGQUIT, SIG_IGN);
    signal(SIGCHLD, parent_handler);
    signal(SIGTTOU, parent_handler);

    /* get tty fd */
    if ((ttyfd = open("/dev/tty", O_RDWR)) < 0) {
        perror("open /dev/tty error");
        exit(0);
    }

    /* save the current dir */
    getcwd(lastdir, MAXLINE);

    init_jobs();
    while(1) {
        set_prompt(prompt);
        cmdline = readline(prompt);
        add_history(cmdline);
        strcpy(buf, cmdline);
        myptr = buf;
        mylim = buf+strlen(buf);
        yylex();
        eval(cmdline);
        reset_argvbuf();
        free(cmdline);
    }
    free_jobs();
    
    return 0;
}
